//! 链路池

use crate::packet;
use alloc::boxed::Box;
use alloc::collections::BTreeMap;
use core::future::pending;
use futures_util::future::{join_all, select_all};
use unmp_link::connection::Connection;

/// 链路描述符
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct Fd(u32);

/// 路由表
pub struct Conns {
    fd_sn: u32,
    conns: BTreeMap<Fd, Connection>,
}
impl Conns {
    pub fn new() -> Self {
        Self {
            fd_sn: 0,
            conns: BTreeMap::new(),
        }
    }
    /// 添加链路到路由表上
    pub fn push(&mut self, conn: Connection) -> Fd {
        log::info!("{:?} connected", conn.as_ref());
        self.fd_sn += 1;
        let fd = Fd(self.fd_sn);
        self.conns.insert(fd, conn);
        fd
    }
}
impl Conns {
    /// 接收数据包
    pub async fn recv(&mut self) -> (Fd, packet::Packet) {
        let conns = &mut self.conns;
        loop {
            if conns.is_empty() {
                pending::<()>().await;
            }
            let recv = conns.iter_mut().map(|(_, conn)| Box::pin(conn.recv()));
            let (result, i, _) = select_all(recv).await;
            let fd = if let Some(fd) = conns.keys().nth(i) {
                *fd
            } else {
                log::warn!("unreachable(recv_data-fd)");
                continue;
            };
            match result {
                Ok(data) => {
                    if let Ok(packet) = packet::decode(data) {
                        return (fd, packet);
                    }
                }
                Err(_) => {
                    let conn = conns.remove(&fd);
                    if let Some(conn) = conn {
                        log::warn!("{:?} disconnected", conn.as_ref());
                    }
                }
            };
        }
    }
    pub async fn send<F: Fn(&Fd) -> bool>(
        &mut self,
        packet: packet::Packet,
        fd_filter: F,
    ) -> Result<(), ()> {
        let data = packet::encode(&packet);
        let send = self.conns.iter_mut().filter_map(|(fd, conn)| {
            if fd_filter(fd) {
                Some(Box::pin(conn.send(data.clone())))
            } else {
                None
            }
        });
        let result = join_all(send).await;
        if result.into_iter().any(|x| x.is_ok()) {
            Ok(())
        } else {
            Err(())
        }
    }
}
